from s4dd import S4forDenovoDesign
from s4dd.torch_callbacks import DenovoDesign  # 必须导入 DenovoDesign
from rdkit import Chem
import torch

# 路径
MODEL_PATH = "./new_models/s4_length_constrained"  # 已训练的模型路径
MAX_LENGTH = 50  # 限制生成序列的最大长度

def main():
    # -------------------------------
    # 1. 加载预训练的S4模型
    # -------------------------------
    s4 = S4forDenovoDesign.from_file(MODEL_PATH)


    # -------------------------------
    # 2. 定义化学有效性奖励函数
    # -------------------------------
    def chemical_validity_reward(smiles_list):
        """
        对分子的化学有效性进行惩罚，使用RDKit验证SMILES字符串是否有效。
        :param smiles_list: SMILES字符串列表
        :return: 化学有效性奖励
        """
        valid_count = 0
        for smiles in smiles_list:
            mol = Chem.MolFromSmiles(smiles)
            if mol is not None:
                valid_count += 1
        return valid_count


    # -------------------------------
    # 3. 定义长度有效性奖励函数
    # -------------------------------
    def length_validity_reward(smiles_list):
        """
        检查分子的长度是否符合要求，长度超过最大长度时给予惩罚。
        :param smiles_list: SMILES字符串列表
        :return: 长度有效性奖励
        """
        length_rewards = []
        for smiles in smiles_list:
            if len(smiles) <= MAX_LENGTH:
                length_rewards.append(1)  # 长度有效的分子给予正奖励
            else:
                length_rewards.append(0)  # 长度超出限制的分子不给奖励
        return sum(length_rewards)


    # -------------------------------
    # 4. 定义综合奖励函数
    # -------------------------------
    def reward_fn(predictions):
        """
        综合奖励函数：根据生成的分子长度和化学有效性来计算奖励。
        """
        # 提取生成的SMILES字符串列表
        smiles_list = [prediction for prediction in predictions]

        # 计算化学有效性奖励
        chemical_rewards = chemical_validity_reward(smiles_list)

        # 计算长度有效性奖励
        length_rewards = length_validity_reward(smiles_list)

        # 综合奖励，给予化学有效性和长度有效性不同的权重
        total_reward = chemical_rewards * 0.7 + length_rewards * 0.3  # 设定化学有效性的权重更高
        return total_reward


    # -------------------------------
    # 5. 定义回调函数 - DenovoDesign
    # -------------------------------
    denovo_callback = DenovoDesign(
        design_fn=lambda t: s4.design_molecules(
            n_designs=32,  # 设计32个新分子
            batch_size=16,  # 每批次大小
            temperature=t  # 控制多样性的温度
        ),
        basedir="./new_models/s4_length_constrained",
        temperatures=[1.0, 1.5, 2.0],  # 设定不同温度下生成的多样性
    )

    # -------------------------------
    # 6. 定义自我迭代优化函数
    # -------------------------------
    def optimize_iterations(iterations=5):
        """
        反复运行训练的程序进行自我迭代优化。
        :param iterations: 迭代次数
        """
        for i in range(iterations):
            print(f"Iteration {i+1}...")

            # 生成分子并进行优化
            designs, lls = s4.design_molecules(
                n_designs=32,       # 设计32个新分子
                batch_size=16,      # 每批次大小
                temperature=1.0     # 控制多样性的温度
            )

            # 输出设计的分子及其log似然值
            print(f"Generated Molecules (Iteration {i+1}):", designs)
            print(f"Log-likelihoods (Iteration {i+1}):", lls)

            # 筛选符合长度和化学有效性要求的分子
            valid_smiles = [smiles for smiles in designs if len(smiles) <= MAX_LENGTH and Chem.MolFromSmiles(smiles) is not None]

            # 如果有效分子数较少，停止迭代
            if len(valid_smiles) == 0:
                print(f"No valid molecules found in iteration {i+1}, stopping early.")
                break

            # 将有效分子作为新的训练数据进行优化
            with open("./datasets/chemblv31/train_optimized.txt", "w") as f:
                for smile in valid_smiles:
                    f.write(f"{smile}\n")

            # 使用生成的有效分子进行训练
            s4.train(
                training_molecules_path="./datasets/chemblv31/train_optimized.txt",  # 训练集路径
                val_molecules_path="./datasets/chemblv31/valid.zip",      # 验证集路径
                callbacks=[denovo_callback]  # 使用生成的有效分子进行优化
            )

            # 保存优化后的模型
            s4.save(f"./new_models/s4_length_constrained_iter_{i+1}")

    # -------------------------------
    # 7. 启动自我迭代优化
    # -------------------------------
    optimize_iterations(iterations=5)  # 设置迭代次数

if __name__ == '__main__':
    main()
